from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langgraph.graph import END, MessageGraph
import inspect
from ..core.base_strategy import BaseStrategy

class StrategyGenerator:
    """策略生成服务核心类"""
    
    def __init__(self, llm):
        self.code_gen_chain = self._create_code_chain(llm)
        self.validator_chain = self._create_validator_chain(llm)
        
    def _create_code_chain(self, llm):
        prompt = PromptTemplate(
            input_variables=["objective", "constraints"],
            template="""
            根据以下投资目标和约束条件生成量化交易策略代码：
            
            投资目标：{objective}
            风险约束：{constraints}
            
            要求：
            1. 继承自BaseStrategy类
            2. 实现initialize和handle_data方法
            3. 使用Python 3.10+语法
            4. 包含类型标注
            """
        )
        return LLMChain(llm=llm, prompt=prompt)
    
    def generate_strategy(self, prompt: dict) -> str:
        """生成策略代码并验证"""
        raw_code = self.code_gen_chain.run(prompt)
        return self._validate_syntax(raw_code)
    
    def _validate_syntax(self, code: str) -> str:
        """语法验证和安全检查"""
        try:
            # 动态编译检测语法错误
            compile(code, '<string>', 'exec')
            # 检查是否包含危险操作
            blacklist = ['os.', 'subprocess', 'eval(']
            if any(b in code for b in blacklist):
                raise SecurityError("检测到不安全代码")
            return code
        except SyntaxError as e:
            return self.validator_chain.run({"error": str(e), "code": code})

class StrategyWorkflow:
    """策略生成工作流（集成langgraph）"""
    
    def __init__(self, llm):
        self.graph = MessageGraph()
        self.graph.add_node("generate", self._generate_node)
        self.graph.add_node("validate", self._validate_node)
        self.graph.add_edge("generate", "validate")
        self.graph.add_edge("validate", END)
        
    def _generate_node(self, state):
        return {"code": generator.generate_strategy(state)}
    
    def _validate_node(self, state):
        try:
            # 动态加载策略类
            namespace = {}
            exec(state['code'], namespace)
            for cls in namespace.values():
                if inspect.isclass(cls) and issubclass(cls, BaseStrategy):
                    return {"status": "success", "class": cls}
            raise ValidationError("未找到有效策略类")
        except Exception as e:
            return {"status": "error", "message": str(e)}